#!/usr/bin/python

from genString import *
from crawlerFile import *
from driver import *
import pipes
import atexit

genSizedStringObj = genSizedString()
genClasses = [genSizedStringObj]

crawlerFileObj = crawlerFile()
crawlers = [crawlerFileObj]

outputDir = "./genOut"
executionDir = "./executing"
finishedExecutionDir = "./executed"

scriptFormat = """
from sulley import *

s_initialize("test")

%s

cmd = ("C:\\Program Files\\Windows Mail\\wab.exe", "/Contact", "FUZZED_INPUT_FILE")

executor = executors.fuzzExecutorFile.fuzzExecuteFile("./fuzz", cmd, timeout=1, save_timeouts=False)
sess                   = sessions.session(executor, session_filename="audits/file_sample.session.%d")
target                 = sessions.target("localhost")

sess.add_target(target)
sess.connect(sess.root, s_get("test"))
sess.fuzz()

"""

if __name__ == '__main__':
    #start crawler
    crawlerFileObj.setSearchTerms(".contact", "./input/")
    print "Items found: ", crawlerFileObj.estimateItemsLeft()
    
    while not crawlerFileObj.isDone():
        inputFilePath = crawlerFileObj.getNext()
        print inputFilePath
        inputFile = open(inputFilePath, 'rb')
        inputBuff = inputFile.read()
        print len(inputBuff)
        fileSize = os.path.getsize(inputFilePath)
        if fileSize > len(inputBuff):
            print "Failed to read entire file into memory"
        inputFile.close()
        
        print "Analyzing %d bytes" %(len(inputBuff))
        createdScripts = genSizedStringObj.writeFuzzScripts(scriptFormat, outputDir, inputBuff)
        
        print "Starting scripts"
        threadCount = 30
        runningThreads = []
        for i in range(0, threadCount): 
            sliceStart = i * (len(createdScripts) / threadCount)
            sliceEnd = (i+1) * (len(createdScripts) / threadCount)
            if sliceEnd > len(createdScripts):
                sliceEnd = len(createdScripts)
            driverThread = driver(outputDir, executionDir, finishedExecutionDir, createdScripts[sliceStart: sliceEnd], i)
            atexit.register(driverThread.shutdown) #kill subprocess if we exit
            driverThread.start()
            runningThreads.append(driverThread)
            
        while len(runningThreads) > 0:    
            try: 
                runningThreads = [t.join(1) for t in runningThreads if t is not None and t.isAlive()]
            except KeyboardInterrupt:
                for t in runningThreads:
                    t.shutdown()


        #send fuzz script to executor 
        